import os
import shutil
import random
from tqdm import tqdm
from PIL import Image
from comic import is_color_comic
from image_utils import split_image

def find_and_copy_color_images(source_dir, dest_dir, limit, per_book_limit=5, per_dir_limit=5, max_height=1568):
    #dest_bw = dest_dir + "_bw"
    #os.makedirs(dest_bw, exist_ok=True)
    os.makedirs(dest_dir, exist_ok=True)
    print(f"彩色图片将被复制到: {dest_dir}")
    print(f"开始在 '{source_dir}' 中搜索，每个漫画目录上限 {per_book_limit} 张，每个章节目录上限 {per_dir_limit} 张，长图阈值 {max_height} ...")

    copied_count = 0
    img_exts = ['.jpg', '.jpeg', '.png', '.webp', '.bmp']
    book_copy_counts = {}
    long_image_processed_dirs = set()

    # 获取所有漫画目录并打乱顺序
    book_dirs = [os.path.join(source_dir, d) for d in os.listdir(source_dir) if os.path.isdir(os.path.join(source_dir, d))]
    random.shuffle(book_dirs)

    with tqdm(desc="已检测图片", unit=" 张") as pbar:
        for book_dir in book_dirs:
            book_name = os.path.basename(book_dir)
            book_count = book_copy_counts.get(book_dir, 0)
            if book_count >= per_book_limit:
                continue

            # 遍历章节目录
            chapter_dirs = [os.path.join(book_dir, d) for d in os.listdir(book_dir) if os.path.isdir(os.path.join(book_dir, d))]
            for chapter_dir in chapter_dirs:
                chapter_name = os.path.basename(chapter_dir)
                chapter_count = 0
                consecutive_bw_count = 0

                files = [f for f in os.listdir(chapter_dir) if any(f.lower().endswith(ext) for ext in img_exts)]
                random.shuffle(files)  # 打乱图片顺序
                for file in files:
                    if copied_count >= limit or book_copy_counts.get(book_dir, 0) >= per_book_limit:
                        break

                    name, ext = os.path.splitext(file)
                    if name.lower().endswith('_color'):
                        continue

                    image_path = os.path.join(chapter_dir, file)
                    pbar.update(1)

                    try:
                        # 先判断是否为彩色图片
                        is_color = is_color_comic(
                            image_path,
                            color_threshold=50,
                            color_ratio_threshold=0.5,
                            min_block_pixels=50000,
                            saturation_threshold=0.25,
                            hist_diff_threshold=0.3,
                            color_block_ratio=1/3
                        )
                        if not is_color:
                            consecutive_bw_count += 1
                            if consecutive_bw_count >= 5:
                                tqdm.write(f"章节 {chapter_name} 连续5张黑白漫画，跳过后续图片。")
                                break
                            continue
                        else:
                            consecutive_bw_count = 0

                        img = Image.open(image_path)
                        width, height = img.size

                        current_book_count = book_copy_counts.get(book_dir, 0)
                        if current_book_count >= per_book_limit:
                            break

                        # 长图处理
                        if height > max_height:
                            # 略超出范围则缩放（保持比例）
                            if height <= max_height * 1.3:
                                scale = max_height / height
                                new_width = int(width * scale)
                                new_img = img.resize((new_width, max_height), Image.LANCZOS)
                                original_filename = os.path.basename(image_path)
                                new_filename = f"{book_name}_{chapter_name}_{original_filename}"
                                dest_file_path = os.path.join(dest_dir, new_filename)
                                new_img.save(dest_file_path)
                                book_copy_counts[book_dir] = book_copy_counts.get(book_dir, 0) + 1
                                copied_count += 1
                                tqdm.write(f"  -> 已缩放图片: {new_filename}")
                                continue
                            # 超出很多则分割
                            if chapter_dir in long_image_processed_dirs:
                                continue
                            long_image_processed_dirs.add(chapter_dir)

                            blocks = split_image(image_path, block_height=max_height, overlap=128)
                            for i, block_img in enumerate(blocks):
                                if copied_count >= limit or book_copy_counts.get(book_dir, 0) >= per_book_limit:
                                    break

                                original_filename_no_ext = os.path.splitext(os.path.basename(image_path))[0]
                                new_filename = f"{book_name}_{chapter_name}_{original_filename_no_ext}_part{i+1}{ext}"
                                dest_file_path = os.path.join(dest_dir, new_filename)
                                # 分割图片
                                block_img.save(dest_file_path)
                                #block_img_bw = block_img.convert('L')
                                #block_img_bw.save(os.path.join(dest_bw, new_filename))

                                book_copy_counts[book_dir] = book_copy_counts.get(book_dir, 0) + 1
                                copied_count += 1
                                tqdm.write(f"  -> 已复制分割块: {new_filename}")
                            continue

                        # 普通彩色图片直接复制
                        original_filename = os.path.basename(image_path)
                        new_filename = f"{book_name}_{chapter_name}_{original_filename}"
                        dest_file_path = os.path.join(dest_dir, new_filename)
                        shutil.copy2(image_path, dest_file_path)

                        # 生成同名黑白图片
                        #img_bw = img.convert('L')
                        #dest_bw_file_path = os.path.join(dest_bw, new_filename)
                        #img_bw.save(dest_bw_file_path)

                        book_copy_counts[book_dir] = book_copy_counts.get(book_dir, 0) + 1
                        copied_count += 1
                        tqdm.write(f"发现彩色图片: {original_filename} -> 已复制为 {new_filename} (漫画: {book_copy_counts[book_dir]}/{per_book_limit}, 总计: {copied_count}/{limit})")

                    except Exception as e:
                        tqdm.write(f"处理图片 {os.path.basename(image_path)} 时出错: {e}")

    print(f"\n搜索完成。总共复制了 {copied_count} 张彩色图片。")

def batch_bw_convert(src_dir, bw_dir):
    os.makedirs(bw_dir, exist_ok=True)
    img_exts = ['.jpg', '.jpeg', '.png', '.webp', '.bmp']
    files = [f for f in os.listdir(src_dir) if any(f.lower().endswith(ext) for ext in img_exts)]
    for file in tqdm(files, desc="黑白化图片", unit="张"):
        src_path = os.path.join(src_dir, file)
        bw_path = os.path.join(bw_dir, file)
        try:
            img = Image.open(src_path)
            img_bw = img.convert('L')
            img_bw.save(bw_path)
        except Exception as e:
            tqdm.write(f"处理图片 {file} 时出错: {e}")

if __name__ == "__main__":
    source = "E:\\software\\jmcomic_v1.2.6_windows_x64\\jmcomic\\cache\\book"
    dest = "E:\\software\\jmcomic_v1.2.6_windows_x64\\jmcomic\\train_color_comics"
    limit = 300
    per_book_limit = 5
    per_dir_limit = 5
    max_height = 1568
    #find_and_copy_color_images(source, dest, limit, per_book_limit, per_dir_limit, max_height)

    #复制完成后批量黑白化
    bw_dest = dest + "_bw"
    batch_bw_convert(dest, bw_dest)
